import os
import glob
import argparse
import gmalg

def verify_keys (prikey_file, pubkey_file):
    sm2 = gmalg.SM2()
    sk = read_binary_file_(prikey_file)
    pk = read_binary_file_(pubkey_file)

    for i in range(1000):
        pk_new  = sm2.generate_pk(sk[0 + i * 32:32 + i * 32])
        assert pk[0 + i * 65:65 + i * 65] == pk_new, "公钥验证失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    args = parser.parse_args()
    directory = args.directory

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    pubkey_file = os.path.join(directory, f'sm2_pubkey')
    prikey_file = os.path.join(directory, f'sm2_prikey')
    verify_keys(prikey_file, pubkey_file)
    # 读取文件内容
    pubkey_content = read_binary_file(pubkey_file)
    prikey_content = read_binary_file(prikey_file)
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(1000):
        # 添加到合并内容
        merged_content += f"公钥=\n{pubkey_content[0 + i * 130:130 + i * 130]}\n"
        merged_content += f"私钥=\n{prikey_content[0 + i * 64:64 + i * 64]}\n"
        merged_content += "\n"
    
    # 写入到txt文件
    output_file = 'sm2密钥对生成速率.txt'
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
